import os,cv2
import numpy as np
from cameraDrv import CameraDriver
import faiss
import time
import keyboard 
import datetime
from invtest import Inv_kinematics,trans2matrix
from ArmDriver import ArmDriver
import pickle
import math

class AcqTask:
    def __init__(self):
        print("")
        self.armctl = ArmDriver('COM5', 115200)
        self.cam = CameraDriver()
        # ret = self.armctl.open_serial()
        # self.cam.openCamera()
        self.X = 189.65
        self.Y = 0.0
        self.Z = 220.6
        self.Rx = 180
        self.Ry = 0
        self.Rz = 180
        self.dis = 1
        self.oldPos_ = [189.65,0.0,220.6]
        self.startFlag_ = False
        self.arucoInitX = 0
        self.arucoInitY = 0
        self.cnt = 0
        self.cap = cv2.VideoCapture(1)
        self.data = {
            "robot0_demo_start_pose": [],
            "robot0_demo_end_pose": [],
            "robot0_eef_pos": [],
            "robot0_eef_rot_axis_angle": [],
            "robot0_gripper_width": [],
            "camera0_rgb": []
        }



    def process_track(self):
        print("###process_track")
        self.data["robot0_demo_start_pose"].append([220.0,0.0,160.0,self.Rx,self.Ry,self.Rz])
        while True:
            retCap, frame = self.cap.read()
            cv2.imshow("frame",frame)
            if cv2.waitKey(1) == ord('q'):
              break

            self.cnt = self.cnt + 1
            color_intrin, depth_intrin, img, img_depth, aligned_depth_frame = self.cam.get_aligned_images()
            imageSrc = img.copy()
            #进行二维码检测，获取Tcm（注：相机坐标系到二维码坐标系）
            Tcm = self.cam.aruc_detect(img)
            if Tcm is None:
                #print("none aruco detect")
                continue
            x = 1000 * Tcm[2,3]
            y = 1000 * Tcm[0,3]
            #print(f"x:{x}  y:{y}")
            if self.startFlag_ == True: 
                if self.arucoInitX < 1:
                    self.arucoInitX = x
                    self.arucoInitY = y
                xBias = self.arucoInitX - x
                yBias = y - self.arucoInitY
                #print(f"xBias:{xBias}  yBias:{yBias}")
                
                #读取机械臂末端实时位姿T06
                T06 = self.armctl.ReadMotorPos()
                #机械臂起始x = 220.0   y = 0.0  
                ctl_x = 220.0 + xBias  #
                ctl_y = 0.0 + yBias
                #maxValue为机械臂的速度控制（偏移越小速度越慢）
                maxValue = 4 * max(math.fabs(ctl_x - T06[0,3]),math.fabs(ctl_y - T06[1,3]))
                if maxValue < 10:
                    maxValue = 10
                if maxValue > 120:
                    maxValue = 120
                self.armctl.defaultSpeed_ = maxValue
                self.armctl.InvControl(ctl_x,ctl_y,160.0,self.Rx,self.Ry,self.Rz)

                ####存储数据####
                #if self.cnt % 2 == 0 and retCap == True:
                if retCap == True:
                    resized_image = cv2.resize(frame, (224,224), interpolation=cv2.INTER_AREA)
                    self.data["robot0_eef_pos"].append([T06[0,3],T06[1,3],160.0])
                    self.data["robot0_eef_rot_axis_angle"].append([self.Rx,self.Ry,self.Rz])
                    self.data["robot0_gripper_width"].append(0)
                    self.data["camera0_rgb"].append(resized_image)

    def processKeyBoard(self,e):

        if e.event_type == keyboard.KEY_DOWN:
            #按下s键，开启采集数据
            if e.name == 's':
                self.startFlag_ = True
                #按下f键，回到起始位置
            elif e.name == 'f':
                self.armctl.defaultSpeed_ = 100
                T06 = self.armctl.ReadMotorPos()
                self.armctl.InvControl(220.0,0.0,220.6,self.Rx,self.Ry,self.Rz)
                self.startFlag_ = False
            #按下g键，保存当前episode文件到本地
            elif e.name == 'g':
                print("######store pklFile success")
                with open('demo50.pkl', 'wb') as file:
                    pickle.dump(self.data, file)


if __name__ == "__main__":

    acqTask = AcqTask()
    _ = acqTask.armctl.open_serial()
    acqTask.cam.openCamera()
    T06 = acqTask.armctl.ReadMotorPos()
    acqTask.armctl.InvControl(220.0,0.0,160.0,acqTask.Rx,acqTask.Ry,acqTask.Rz)
    time.sleep(1.0)
    keyboard.hook(acqTask.processKeyBoard)
    acqTask.process_track()